from __future__ import annotations
import os
# 设置逻辑核心数（可根据你的电脑配置调整，如4、8等，建议填写实际逻辑核心数）
os.environ["LOKY_MAX_CPU_COUNT"] = "4"
import pickle
from typing import List

import matplotlib.pyplot as plt
import numpy as np

from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score

try:
	# optional progress bar
	from tqdm import tqdm
except Exception:
	# fallback if tqdm not available: implement a simple textual progress generator
	def tqdm(iterable, desc: str | None = None, unit: str | None = None):
		# Try to compute total length for better display
		total = len(iterable) if hasattr(iterable, "__len__") else None
		idx = 0
		bar_width = 30
		prefix = (desc + ' ' if desc else '') + (f"({unit}) " if unit else '')
		for item in iterable:
			idx += 1
			if total:
				pct = idx / total
				filled = int(bar_width * pct)
				bar = "#" * filled + "-" * (bar_width - filled)
				print(f"\r{prefix}[{bar}] {idx}/{total}", end="", flush=True)
			else:
				print(f"\r{prefix}{idx}", end="", flush=True)
			yield item
		# final newline after iteration
		print()


def main() -> None:
	# 加载数字数据集
	digits = load_digits()
	X = digits.data
	y = digits.target

	# 将数据集划分为训练集和测试集
	X_train, X_test, y_train, y_test = train_test_split(
		X, y, test_size=0.2, random_state=42, stratify=y
	)

	# 初始化变量以存储最佳准确率，相应的k值和最佳knn模型
	best_accuracy = 0.0
	best_k = 1
	best_model = None

	# 初始化一个列表以存储每个k值的准确率
	ks: List[int] = list(range(1, 41))
	accuracies: List[float] = []

	# 尝试从1到40的k值，对于每个k值，训练knn模型，保存最佳准确率，k值和knn模型
	for k in tqdm(ks, desc="提示词:", unit="每秒的单位"):
		clf = KNeighborsClassifier(n_neighbors=k)
		clf.fit(X_train, y_train)
		preds = clf.predict(X_test)
		acc = accuracy_score(y_test, preds)
		accuracies.append(acc)

		if acc > best_accuracy:
			best_accuracy = acc
			best_k = k
			best_model = clf

	# 将最佳KNN模型保存到二进制文件
	if best_model is not None:
		try:
			with open("best_knn_model.pkl", "wb") as f:
				pickle.dump(best_model, f)
		except PermissionError:
			# If the file is open or not writable, save to a fallback name and inform the user
			fallback_model = "best_knn_model_fallback.pkl"
			with open(fallback_model, "wb") as f:
				pickle.dump(best_model, f)
			print(
				f"Could not write 'best_knn_model.pkl' (permission denied). Saved model to '{fallback_model}' instead."
			)

	# 绘制准确率随k变化的折线图并保存为 PDF
	plt.figure(figsize=(8, 6))
	plt.plot(ks, accuracies, marker="o")
	plt.xlabel("k value")
	plt.ylabel("Accuracy")
	plt.title("Accuracy of different k values")

	# 绘制垂直线和标注最优 k
	plt.axvline(best_k, color="red", linestyle="-", linewidth=1.5)
	# annotate with the best accuracy
	plt.text(
		best_k,
		max(accuracies),
		f"k={best_k}, Accuracy={best_accuracy:.2f}",
		color="red",
		fontsize=10,
		horizontalalignment="left",
		verticalalignment="bottom",
	)

	plt.grid(True, linestyle="--", alpha=0.4)
	plt.tight_layout()
	# Prefer to save atomically: write to a temp file then replace the target filename.
	target_pdf = "accuracy_plot.pdf"
	temp_pdf = "accuracy_plot.tmp.pdf"
	try:
		plt.savefig(temp_pdf, format="pdf", bbox_inches="tight")
		# try to atomically replace the target; retry if locked
		max_retries = 5
		for attempt in range(1, max_retries + 1):
			try:
				os.replace(temp_pdf, target_pdf)
				break
			except PermissionError:
				# target may be open/locked by a viewer; wait a bit and retry
				import time

				if attempt == max_retries:
					raise
				time.sleep(0.5)
		else:
			# If loop completed without break, raise a generic error
			raise PermissionError(f"Could not replace {target_pdf} after {max_retries} attempts")
	except PermissionError as e:
		# If we couldn't write or replace, inform user and keep the temp PDF name
		# Attempt a best-effort: if temp_pdf exists but couldn't be moved, rename to a fallback readable name
		fallback_pdf = "accuracy_plot_fallback.pdf"
		try:
			if os.path.exists(temp_pdf):
				os.replace(temp_pdf, fallback_pdf)
			print(
				f"Could not write or replace '{target_pdf}' (permission denied). Saved plot to '{fallback_pdf}' instead."
			)
		except Exception:
			# If even fallback fails, raise the original error
			raise e
	plt.close()

	# 打印最佳准确率和相应的k值
	print(f"Best k: {best_k}, Best accuracy: {best_accuracy:.4f}")


if __name__ == "__main__":
	main()